/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;
import java.io.*;
import java.text.*;

import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.TrellisSlice.SliceItem;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class Trellis {
	private ArrayList<TrellisSlice> slices;
	private int currentTime;
	private int debug;
	
	public Trellis() {
		slices = new ArrayList<TrellisSlice>(100);
		currentTime = 0;
	}
	
	public void newSlice(long overtFactors) {
		currentTime++;
		TrellisSlice slice = new TrellisSlice(overtFactors, this, currentTime);
		slices.add(slice);
		//slice.setCumulativeLogScale(slice.getPrevious().getCumulativeLogScale() + slice.getPrevious().getSliceLogScale());
		// set cumulative scale
		slice.scale(1.0);
		
		if (debug > 1) {
			System.out.printf("New Slice #%d (%s)\n", currentTime, FactorTuple.toStringNoNull(overtFactors));
		}
	}
	
	public TrellisSlice getSliceByOffset(int timeOffset) {
		assert(timeOffset <= 0 && timeOffset + currentTime >= 0);
		return slices.get(currentTime - 1 + timeOffset);
	}
	
	protected TrellisSlice getPreviousSlice(TrellisSlice slice) {
		return slices.get(slice.getTime() - 1 - 1);
	}
	
	protected TrellisSlice getNextSlice(TrellisSlice slice) {
		if (slices.size() <= slice.getTime()) return null;
		return slices.get(slice.getTime() - 1 + 1);
	}
	
	public TrellisSlice getCurrentSlice() {
		return getSliceByOffset(0);
	}
	
	public void dumpTrellis(String filename) {
		dumpTrellis(filename, true);
	}
	
	public void dumpTrellis(String filename, boolean includeDeadEnd) {
		try {
			dumpTrellis(new FileOutputStream(filename), includeDeadEnd);
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		}
	}
	
	public void dumpTrellis(OutputStream os) {
		dumpTrellis(os, true);
	}
	
	public void dumpTrellis(OutputStream os, boolean includeDeadEnd) {
		OutputStreamWriter out = new OutputStreamWriter(os);
		DecimalFormat nf = new DecimalFormat("0.###E0");
		nf.setMaximumFractionDigits(3);
		
		HashSet<SliceItem> deadEnds = new HashSet<SliceItem>();
		try {
			out.append("digraph Trellis {\n");
			out.append("  rankdir=LR\n;");
			for(TrellisSlice slice : slices) {
				deadEnds.clear();
				for(SliceItem item : slice.getElements()) {
					if (!includeDeadEnd && slice != getCurrentSlice()) {
						if (slice.getNext().prev2curr.get(item) == null) {
							// drop items which have to continuation in the next slice
							deadEnds.add(item);
							continue;
						}
					}
					out.append("  N" + item.toString());
					out.append(" [label=\"N" + item.toString());
					if (!item.clusterid.isEmpty()) {
						out.append(" (" + item.clusterid + ")");
					}
					try {
						String prefix = item.tree.getPrefix().toString();
						if (!prefix.isEmpty()) {
							out.append("\\nprefix=");
							out.append(prefix);
						}
					} catch(UnsupportedOperationException e) {
						// item.tree.getPrefix() will throw this if it's a ghost
					}
					
					out.append("\\nprev=" + nf.format(item.getPrevProb() * Math.pow(10.0, slice.getCumulativeLogScale() - slice.getSliceLogScale())));
					out.append("\\nprob=" + nf.format(item.tree.getTotalProb() * Math.pow(10.0, slice.getSliceLogScale())) + "\"];\n");
				}
				out.append(" {rank=same;");
				for(SliceItem item : slice.getElements()) {
					if (includeDeadEnd || !deadEnds.contains(item)) {
						out.append("  N" + item.toString());
					}
				}
				out.append("};\n");
				
				for(Map.Entry<SliceItem,ArrayList<SliceItem>> prevs : slice.curr2prev.entrySet()) {
					if (!includeDeadEnd && deadEnds.contains(prevs.getKey())) {
						continue;
					}
					ArrayList<SliceItem> prevLinks = prevs.getValue();
					if (prevLinks != null) {
						out.append("N" + prevs.getKey().toString() + " -> {");
						for(SliceItem prev : prevLinks) {
							out.append(" N" + prev.toString());
						}
						out.append("};\n");
					}
				}
				/*
				for(Map.Entry<SliceItem,ArrayList<SliceItem>> currs : slice.prev2curr.entrySet()) {
					out.append("N" + currs.getKey().toString() + " -> {");
					for(SliceItem prev : currs.getValue()) {
						out.append(" N" + prev.toString());
					}
					out.append("};\n");
				}
				*/
			}
			out.append("}\n");
			out.flush();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public int getDebug() {
		return debug;
	}

	public void setDebug(int debug) {
		this.debug = debug;
	}

	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		boolean first = true;
		for(TrellisSlice slice : slices) {
			if (!first) {
				sb.append(' ');
			}
			first = false;
			sb.append(FactorTuple.toStringNoNull(slice.getOvertFactors()));
		}
		return sb.toString();
	}
	
	
}
